package com.ds.fans.liteflow.test;

import java.util.*;

public class DAGShortestPath {
    // 存储图的邻接表
    private Map<Integer, List<Edge>> graph;

    // 边类
    static class Edge {
        int destination;
        int weight;

        Edge(int destination, int weight) {
            this.destination = destination;
            this.weight = weight;
        }
    }

    public DAGShortestPath() {
        graph = new HashMap<>();
    }

    // 添加边
    public void addEdge(int source, int destination, int weight) {
        if (!graph.containsKey(source)) {
            graph.put(source, new ArrayList<>());
        }
        graph.get(source).add(new Edge(destination, weight));
    }

    // 拓扑排序
    public List<Integer> topologicalSort() {
        Map<Integer, Integer> inDegree = new HashMap<>();
        Queue<Integer> queue = new LinkedList<>();
        List<Integer> result = new ArrayList<>();
// 初始化入度
        for (Integer vertex : graph.keySet()) {
            inDegree.put(vertex, 0);
        }
        for (List<Edge> edges : graph.values()) {
            for (Edge edge : edges) {
                inDegree.put(edge.destination, inDegree.getOrDefault(edge.destination, 0) + 1);
            }
        }
// 将入度为 0 的顶点入队
        for (Integer vertex : graph.keySet()) {
            if (inDegree.get(vertex) == 0) {
                queue.add(vertex);
            }
        }
        while (!queue.isEmpty()) {
            int vertex = queue.poll();
            result.add(vertex);
            if(graph!=null&&!graph.isEmpty()&&graph.get(vertex)!=null){
                for (Edge edge : graph.get(vertex)) {
                    int destination = edge.destination;
                    inDegree.put(destination, inDegree.get(destination) - 1);
                    if (inDegree.get(destination) == 0) {
                        queue.add(destination);
                    }
                }
            }
        }
        return result;
    }

    // 计算最短路径
    public Map<Integer, Integer> shortestPath(int source) {
        List<Integer> topologicalOrder = topologicalSort();
        Map<Integer, Integer> distance = new HashMap<>();
        for (Integer vertex : graph.keySet()) {
            distance.put(vertex, Integer.MAX_VALUE);
        }
        distance.put(source, 0);
        for (Integer vertex : topologicalOrder) {
            if (distance.get(vertex)!=null&&distance.get(vertex) != Integer.MAX_VALUE) {
                for (Edge edge : graph.get(vertex)) {
                    int destination = edge.destination;
                    int newDistance = distance.get(vertex) + edge.weight;
                    if (distance.get(destination)!=null&&newDistance < distance.get(destination)) {
                        distance.put(destination, newDistance);
                    }
                }
            }
        }
        return distance;
    }

    public static void main(String[] args) {
        DAGShortestPath dag = new DAGShortestPath();
        dag.addEdge(0, 1, 5);
        dag.addEdge(0, 2, 3);
        dag.addEdge(1, 3, 6);
        dag.addEdge(1, 2, 2);
        dag.addEdge(2, 4, 4);
        dag.addEdge(2, 3, 7);
        dag.addEdge(3, 4, -1);
        int source = 0;
        Map<Integer, Integer> shortestDistances = dag.shortestPath(source);
        for (Map.Entry<Integer, Integer> entry : shortestDistances.entrySet()) {
            System.out.println("从顶点 " + source + " 到顶点 " + entry.getKey() + " 的最短距离为: " + entry.getValue());
        }
    }
}

